import sys
import cv2
import time
from PyQt5.QtWidgets import (
    QApplication, QMainWindow, QPushButton,
    QVBoxLayout, QHBoxLayout, QWidget, QFrame,
    QGraphicsView, QGraphicsScene, QGraphicsPixmapItem, QLabel
)
from PyQt5.QtGui import QPixmap, QImage, QFont, QPainter
from PyQt5.QtCore import Qt, QThread, pyqtSignal
from ultralytics import YOLO
from multiprocessing import Queue
from pathlib import Path
import json
from threading import Thread


def load_config():
    """加载配置文件"""
    if getattr(sys, 'frozen', False):
        base_dir = Path(sys._MEIPASS)
    else:
        base_dir = Path(__file__).parent

    config_path = base_dir / "config.json"
    try:
        with open(config_path, "r", encoding="utf-8") as f:
            config = json.load(f)
    except FileNotFoundError:
        sys.exit(f"错误：配置文件 {config_path} 不存在")
    except json.JSONDecodeError:
        sys.exit(f"错误：配置文件格式不正确")

    model_path = Path(config["model_path"])
    if not model_path.exists():
        sys.exit(f"错误：模型文件 {model_path} 不存在")

    return {
        "video_source": config["video_source"],
        "model_path": str(model_path),
        "save_dirs": config["save_dirs"],
        "alert_classes": config.get("alert_classes", ["nohel"]),
    }


# 保存原始帧（异步）
class RawFrameSaver:
    def __init__(self, save_dir="raw_frames"):
        config = load_config()
        self.save_dir = Path(config["save_dirs"]["raw_frames"])
        self.save_dir.mkdir(exist_ok=True)
        self.frame_queue = Queue(maxsize=20)
        self.thread = Thread(target=self._save_worker, daemon=True)
        self.thread.start()

    def add_frame(self, frame, timestamp_str):
        if not self.frame_queue.full():
            self.frame_queue.put((frame.copy(), timestamp_str))

    def _save_worker(self):
        while True:
            if self.frame_queue.empty():
                time.sleep(0.05)
                continue
            frame, timestamp_str = self.frame_queue.get()
            cv2.imwrite(str(self.save_dir / f"{timestamp_str}.jpg"), frame)


class FrameSaver:
    def __init__(self, save_dir="full_frames"):
        config = load_config()
        self.save_dir = Path(config["save_dirs"]["full_frames"])
        self.save_dir.mkdir(exist_ok=True)
        self.frame_queue = Queue(maxsize=20)
        self.thread = Thread(target=self._save_worker, daemon=True)
        self.thread.start()

    def save_full_frame(self, frame, timestamp_str):
        if not self.frame_queue.full():
            self.frame_queue.put((frame.copy(), timestamp_str))

    def _save_worker(self):
        while True:
            if self.frame_queue.empty():
                time.sleep(0.05)
                continue
            frame, timestamp_str = self.frame_queue.get()
            cv2.imwrite(str(self.save_dir / f"{timestamp_str}.jpg"), frame)


class ImageSaver:
    def __init__(self, save_dir="detected_objects"):
        config = load_config()
        self.save_dir = Path(config["save_dirs"]["detected_objects"])
        self.save_dir.mkdir(exist_ok=True)
        self.min_interval = 0.3
        self.last_save_time = {}

    def save_cropped_image(self, frame, box, class_name, conf, track_id):
        current_time = time.time()
        if track_id in self.last_save_time and current_time - self.last_save_time[track_id] < self.min_interval:
            return
        self.last_save_time[track_id] = current_time

        x1, y1, x2, y2 = map(int, box)
        cropped = frame[y1:y2, x1:x2]
        if cropped.size == 0:
            return

        timestamp = time.strftime("%Y%m%d_%H%M%S")
        filename = f"{timestamp}_{class_name}_{conf:.2f}_id{track_id}.jpg"
        save_path = self.save_dir / filename

        try:
            ret, buffer = cv2.imencode('.jpg', cropped, [cv2.IMWRITE_JPEG_QUALITY, 90])
            if ret:
                buffer.tofile(str(save_path))
        except Exception as e:
            print(f"保存失败: {str(e)}")


# 推理处理器
class InferenceProcessor:
    def __init__(self, model_path, output_video_dir="output"):
        self.model = YOLO(model_path)
        self.saver = ImageSaver()
        self.frame_saver = FrameSaver()
        self.raw_saver = RawFrameSaver()
        self.class_colors = {
            "hel": (0, 255, 0),
            "nohel": (0, 0, 255),
            "wheel": (255, 0, 0),
        }
        config = load_config()
        self.alert_classes = config["alert_classes"]
        self.alerted_tracks = set()
        self.alert_count = 0

        self.last_detections = []

        # 视频保存
        self.output_video_dir = Path(output_video_dir)
        self.output_video_dir.mkdir(exist_ok=True)
        self.video_writer = None
        self.fourcc = cv2.VideoWriter_fourcc(*'mp4v')
        self.output_video_path = None

        self.original_fps = 25.0
        self.frame_width = 640
        self.frame_height = 480

    def get_unique_video_path(self):
        timestamp = time.strftime("%Y%m%d_%H%M%S")
        return self.output_video_dir / f"detected_{timestamp}.mp4"

    def init_video_writer(self, width, height, fps):
        if self.video_writer is None:
            self.frame_width = width
            self.frame_height = height
            self.original_fps = fps
            self.output_video_path = self.get_unique_video_path()

            self.video_writer = cv2.VideoWriter(
                str(self.output_video_path),
                self.fourcc,
                fps,
                (width, height)
            )
            print(f"📹 开始录制检测视频: {self.output_video_path} | {width}x{height} @ {fps:.1f}fps")

    def process_frame(self, frame):
        raw_frame = frame.copy()
        timestamp_ms = int(time.time() * 1000)

        self.last_detections.clear()

        results = self.model.track(
            frame,
            imgsz=320,
            conf=0.5,
            iou=0.45,
            persist=True,
            device='cpu',       # 使用 CPU 避免 CUDA 冲突
            tracker="bytetrack.yaml",
            verbose=False
        )

        for result in results:
            if result.boxes is None or result.boxes.id is None:
                continue

            boxes = result.boxes.xyxy.cpu().numpy()
            track_ids = result.boxes.id.cpu().numpy().astype(int)
            class_ids = result.boxes.cls.cpu().numpy().astype(int)
            confidences = result.boxes.conf.cpu().numpy()

            for box, track_id, cls_id, conf in zip(boxes, track_ids, class_ids, confidences):
                class_name = result.names[cls_id]
                x1, y1, x2, y2 = map(int, box)
                color = self.class_colors.get(class_name, (255, 255, 255))

                cv2.rectangle(frame, (x1, y1), (x2, y2), color, 2)
                cv2.putText(frame, f"{class_name} ID:{track_id}",
                            (x1, max(y1 - 10, 15)),
                            cv2.FONT_HERSHEY_SIMPLEX, 0.6, (255, 255, 255), 1, cv2.LINE_AA)

                self.last_detections.append((x1, y1, x2, y2, track_id, class_name, conf))

                if class_name in self.alert_classes:
                    obj_key = (class_name, track_id)
                    if obj_key not in self.alerted_tracks:
                        print(f"🚨 报警：{class_name} ID:{track_id}")
                        self.saver.save_cropped_image(frame, [x1, y1, x2, y2], class_name, conf, track_id)
                        self.raw_saver.add_frame(raw_frame, f"alert_{timestamp_ms}_id{track_id}")
                        self.frame_saver.save_full_frame(frame.copy(), f"alert_{timestamp_ms}_id{track_id}")
                        self.alerted_tracks.add(obj_key)
                        self.alert_count += 1

        return frame

    def draw_cached_boxes(self, frame):
        for (x1, y1, x2, y2, track_id, class_name, conf) in self.last_detections:
            color = self.class_colors.get(class_name, (255, 255, 255))
            cv2.rectangle(frame, (x1, y1), (x2, y2), color, 2)
            cv2.putText(frame, f"{class_name} ID:{track_id}",
                        (x1, max(y1 - 10, 15)),
                        cv2.FONT_HERSHEY_SIMPLEX, 0.6, (255, 255, 255), 1, cv2.LINE_AA)
        return frame

    def release(self):
        if self.video_writer:
            self.video_writer.release()
            print(f"✅ 检测视频已保存: {self.output_video_path}")
            self.video_writer = None


# 视频处理线程
class VideoThread(QThread):
    frame_ready = pyqtSignal(QImage)
    status_update = pyqtSignal(str)
    stats_update = pyqtSignal(int, float)

    def __init__(self, processor, video_source):
        super().__init__()
        self.processor = processor
        self.video_source = video_source
        self.running = True
        self.frame_skip = 1
        self.frame_counter = 0
        self.cap = None

    def run(self):
        self.cap = cv2.VideoCapture(self.video_source)
        if not self.cap.isOpened():
            self.status_update.emit("❌ 无法打开视频源")
            return

        fps = self.cap.get(cv2.CAP_PROP_FPS)
        width = int(self.cap.get(cv2.CAP_PROP_FRAME_WIDTH))
        height = int(self.cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
        if fps <= 0 or fps > 100:
            fps = 25.0

        self.processor.init_video_writer(width, height, fps)

        MAX_DISPLAY_W, MAX_DISPLAY_H = 1600, 900
        display_interval = 1 / 25
        last_send_time = time.time()

        frame_count = 0
        start_time = time.time()
        processed_frame_count = 0

        while self.running:
            ret, frame = self.cap.read()
            if not ret:
                self.status_update.emit(f"⏹️ 视频结束 ({processed_frame_count}帧)")
                break

            self.frame_counter += 1

            if self.frame_counter % (self.frame_skip + 1) == 0:
                processed = self.processor.process_frame(frame)
                processed_frame_count += 1
            else:
                processed = self.processor.draw_cached_boxes(frame.copy())

            self.processor.video_writer.write(processed)

            current_time = time.time()
            if current_time - last_send_time < display_interval:
                continue
            last_send_time = current_time

            h, w = processed.shape[:2]

            # 安全缩放
            scale = min(MAX_DISPLAY_W / w, MAX_DISPLAY_H / h, 1.0)
            new_w = max(1, int(w * scale))
            new_h = max(1, int(h * scale))

            if scale < 1:
                display_img = cv2.resize(processed, (new_w, new_h), interpolation=cv2.INTER_AREA)
            else:
                display_img = processed.copy()

            rgb = cv2.cvtColor(display_img, cv2.COLOR_BGR2RGB).copy()
            h, w, ch = rgb.shape
            bytes_per_line = ch * w

            try:
                image_bytes = rgb.tobytes()
                qimg = QImage(image_bytes, w, h, bytes_per_line, QImage.Format_RGB888)
                qimg.__image_bytes_ref = image_bytes  # 延长生命周期
                self.frame_ready.emit(qimg)
            except Exception as e:
                print(f"❌ 图像转换失败: {e}")

            frame_count += 1
            elapsed = time.time() - start_time
            if elapsed >= 1.0:
                real_fps = frame_count / elapsed
                alert_count = self.processor.alert_count
                self.stats_update.emit(alert_count, real_fps)
                frame_count = 0
                start_time = time.time()

        self.cap.release()

    def stop(self):
        self.running = False
        self.wait()


# 主窗口
class MainWindow(QMainWindow):
    def __init__(self):
        super().__init__()
        self.setWindowTitle("🛡️ YOLOv8 安全检测系统")
        self.setGeometry(100, 100, 1600, 900)
        self.setStyleSheet("background-color: #f0f0f0; font-family: Arial;")
        # --- ✅ 禁用最大化按钮 ---
        self.setWindowFlags(
            Qt.Window |  # 基础窗口
            Qt.WindowCloseButtonHint |  # 保留关闭按钮
            Qt.WindowMinimizeButtonHint  # 保留最小化按钮
            # ❌ 没有 Qt.WindowMaximizeButtonHint → 禁用最大化
        )

        # 可选：禁止窗口拉伸（固定大小）
        self.setFixedSize(1600, 900)

        config = load_config()
        for dir_path in config["save_dirs"].values():
            Path(dir_path).mkdir(exist_ok=True)

        self.processor = InferenceProcessor(
            model_path=config["model_path"],
            output_video_dir="output"
        )

        self.video_thread = None
        self.setup_ui()
        self.center_window()

    def setup_ui(self):
        central_widget = QWidget()
        self.setCentralWidget(central_widget)

        layout = QHBoxLayout(central_widget)
        layout.setContentsMargins(15, 15, 15, 15)
        layout.setSpacing(15)

        # 左侧：使用 QGraphicsView 显示（安全）
        self.graphics_view = QGraphicsView()
        self.graphics_scene = QGraphicsScene()
        self.graphics_view.setScene(self.graphics_scene)
        self.graphics_view.setRenderHint(QPainter.Antialiasing)
        self.graphics_view.setRenderHint(QPainter.SmoothPixmapTransform)
        self.graphics_view.setAlignment(Qt.AlignCenter)
        self.graphics_view.setStyleSheet("background-color: #000; border: none;")

        self.pixmap_item = QGraphicsPixmapItem()
        self.graphics_scene.addItem(self.pixmap_item)

        # 右侧：控制面板
        control_frame = QFrame()
        control_frame.setFrameShape(QFrame.StyledPanel)
        control_frame.setStyleSheet("""
            background-color: white;
            border-radius: 12px;
            padding: 15px;
            border: 1px solid #ddd;
        """)
        control_layout = QVBoxLayout(control_frame)

        title = QLabel("控制面板")
        title.setFont(QFont("Arial", 16, QFont.Bold))
        title.setAlignment(Qt.AlignCenter)
        control_layout.addWidget(title)

        self.start_btn = QPushButton("▶️ 开始检测")
        self.start_btn.setStyleSheet(self.btn_style("#4CAF50"))
        self.start_btn.clicked.connect(self.start_detection)

        self.pause_btn = QPushButton("⏸️ 暂停")
        self.pause_btn.setStyleSheet(self.btn_style("#FF9800"))
        self.pause_btn.clicked.connect(self.pause_detection)
        self.pause_btn.setEnabled(False)

        self.exit_btn = QPushButton("⏹️ 退出")
        self.exit_btn.setStyleSheet(self.btn_style("#F44336"))
        self.exit_btn.clicked.connect(self.close)

        for btn in [self.start_btn, self.pause_btn, self.exit_btn]:
            btn.setFixedHeight(50)
            control_layout.addWidget(btn)

        self.status_label = QLabel("状态：就绪")
        self.alert_label = QLabel("报警数：0")
        self.fps_label = QLabel("FPS：0")

        for label in [self.status_label, self.alert_label, self.fps_label]:
            label.setFont(QFont("Arial", 11))
            label.setStyleSheet("padding: 5px;")
            control_layout.addWidget(label)

        control_layout.addStretch()

        layout.addWidget(self.graphics_view, 3)
        layout.addWidget(control_frame, 1)

    def btn_style(self, color):
        return f"padding: 12px; background-color: {color}; color: white; border: none; border-radius: 8px;"

    def center_window(self):
        qr = self.frameGeometry()
        cp = QApplication.desktop().availableGeometry().center()
        qr.moveCenter(cp)
        self.move(qr.topLeft())

    def start_detection(self):
        if self.video_thread and self.video_thread.isRunning():
            return

        config = load_config()
        self.video_thread = VideoThread(self.processor, config["video_source"])
        self.video_thread.frame_ready.connect(self.update_frame)
        self.video_thread.status_update.connect(self.update_status)
        self.video_thread.stats_update.connect(self.update_stats)
        self.video_thread.start()

        self.start_btn.setEnabled(False)
        self.pause_btn.setEnabled(True)
        self.status_label.setText("状态：检测中...")

    def pause_detection(self):
        if self.video_thread:
            self.video_thread.stop()
            self.video_thread = None
        self.start_btn.setEnabled(True)
        self.pause_btn.setEnabled(False)
        self.status_label.setText("状态：已暂停")

    def update_frame(self, qimg):
        try:
            if qimg and not qimg.isNull():
                pixmap = QPixmap.fromImage(qimg)
                if not pixmap.isNull():
                    self.pixmap_item.setPixmap(pixmap)
                    self.graphics_view.fitInView(self.pixmap_item, Qt.KeepAspectRatio)
            else:
                print("❌ 接收到空图像")
        except Exception as e:
            print(f"❌ 更新画面失败: {e}")

    def update_stats(self, alert_count, fps):
        self.alert_label.setText(f"报警数：{alert_count}")
        self.fps_label.setText(f"FPS：{fps:.1f}")

    def update_status(self, text):
        self.status_label.setText(f"状态：{text}")

    def closeEvent(self, event):
        if self.video_thread and self.video_thread.isRunning():
            self.video_thread.stop()
        self.processor.release()
        event.accept()


if __name__ == "__main__":
    app = QApplication(sys.argv)
    window = MainWindow()
    window.showMaximized()  # ✅ 可安全最大化
    sys.exit(app.exec_())